package license

import (
    "crypto"
    "crypto/aes"
    "crypto/cipher"
    "crypto/rand"
    "crypto/rsa"
    "crypto/sha256"
    "crypto/x509"
    "encoding/base64"
    "encoding/hex"
    "encoding/json"
    "encoding/pem"
    "fmt"
    "io/ioutil"
    "os"
    "path/filepath"
    "strings"
    "time"
    "hardwareInfo/hardware" // 添加对hardware包的导入
)

// License 许可证结构
type License struct {
    HardwareHash    string `json:"hardwareHash"`
    EncryptedExpiry string `json:"encryptedExpiry"` // 加密的过期时间和管理员信息
    Signature       string `json:"signature"`
    // 移除明文admin字段，admin信息将加密存储在EncryptedExpiry中
}

// SensitiveData 敏感数据结构（用于加密存储）
type SensitiveData struct {
    ExpiryDate string `json:"expiryDate"`
    Admin      bool   `json:"admin"`
}

// LicenseStatus 许可证状态
type LicenseStatus struct {
    Valid  bool   `json:"valid"`
    Reason string `json:"reason"`
}

// getKeyFilePath 获取密钥文件路径，支持从可执行文件目录和工作目录查找
func getKeyFilePath(filename string) (string, error) {
    // 获取可执行文件所在目录
    ex, err := os.Executable()
    if err != nil {
        return "", fmt.Errorf("获取可执行文件路径失败: %v", err)
    }
    exPath := filepath.Dir(ex)
    
    // 首先尝试可执行文件目录
    keyPath := filepath.Join(exPath, "config", filename)
    if _, err := os.Stat(keyPath); err == nil {
        return keyPath, nil
    }
    
    // 如果找不到，则尝试当前工作目录
    keyPath = filepath.Join("config", filename)
    if _, err := os.Stat(keyPath); err == nil {
        return keyPath, nil
    }
    
    return "", fmt.Errorf("未找到密钥文件: %s", filename)
}

// calculateSHA256Hash 计算字符串的SHA256哈希值并返回十六进制字符串
func calculateSHA256Hash(data string) string {
    hasher := sha256.New()
    hasher.Write([]byte(data))
    hashed := hasher.Sum(nil)
    
    // 转换为十六进制字符串
    var sb strings.Builder
    for _, b := range hashed {
        sb.WriteString(fmt.Sprintf("%02x", b))
    }
    
    return sb.String()
}

// encryptLicenseData 使用AES加密许可证敏感信息（过期时间和管理员标识）
func encryptLicenseData(expiryDate string, admin bool, hardwareHash string) (string, error) {
    // 构造包含过期时间和管理员标识的数据
    sensitiveData := SensitiveData{
        ExpiryDate: expiryDate,
        Admin:      admin,
    }
    
    // 使用硬件哈希的前32位作为密钥
    keyData := []byte(hardwareHash[:32])
    
    // 创建AES块加密器
    block, err := aes.NewCipher(keyData)
    if err != nil {
        return "", fmt.Errorf("创建AES加密器失败: %v", err)
    }
    
    // 序列化数据
    dataStr, err := json.Marshal(sensitiveData)
    if err != nil {
        return "", fmt.Errorf("序列化敏感数据失败: %v", err)
    }
    
    // 生成随机IV
    iv := make([]byte, aes.BlockSize)
    if _, err := rand.Read(iv); err != nil {
        return "", fmt.Errorf("生成IV失败: %v", err)
    }
    
    // 填充数据
    paddedData := pkcs7Pad(dataStr, aes.BlockSize)
    
    // 加密
    mode := cipher.NewCBCEncrypter(block, iv)
    encrypted := make([]byte, len(paddedData))
    mode.CryptBlocks(encrypted, paddedData)
    
    // 将IV和加密数据组合
    result := hex.EncodeToString(iv) + ":" + hex.EncodeToString(encrypted)
    return result, nil
}

// decryptLicenseData 使用AES解密许可证敏感信息
func decryptLicenseData(encryptedData string, hardwareHash string) (*SensitiveData, error) {
    // 分割IV和加密数据
    parts := strings.Split(encryptedData, ":")
    if len(parts) != 2 {
        return nil, fmt.Errorf("无效的加密格式")
    }
    
    iv, err := hex.DecodeString(parts[0])
    if err != nil {
        return nil, fmt.Errorf("解码IV失败: %v", err)
    }
    
    encrypted, err := hex.DecodeString(parts[1])
    if err != nil {
        return nil, fmt.Errorf("解码加密数据失败: %v", err)
    }
    
    // 使用硬件哈希的前32位作为密钥
    keyData := []byte(hardwareHash[:32])
    
    // 创建AES块解密器
    block, err := aes.NewCipher(keyData)
    if err != nil {
        return nil, fmt.Errorf("创建AES解密器失败: %v", err)
    }
    
    // 解密
    mode := cipher.NewCBCDecrypter(block, iv)
    decrypted := make([]byte, len(encrypted))
    mode.CryptBlocks(decrypted, encrypted)
    
    // 去除填充
    unpaddedData, err := pkcs7Unpad(decrypted)
    if err != nil {
        return nil, fmt.Errorf("去除填充失败: %v", err)
    }
    
    // 反序列化数据
    var sensitiveData SensitiveData
    if err := json.Unmarshal(unpaddedData, &sensitiveData); err != nil {
        return nil, fmt.Errorf("反序列化敏感数据失败: %v", err)
    }
    
    return &sensitiveData, nil
}

// pkcs7Pad 进行PKCS7填充
func pkcs7Pad(data []byte, blockSize int) []byte {
    padding := blockSize - len(data)%blockSize
    padText := make([]byte, padding)
    for i := range padText {
        padText[i] = byte(padding)
    }
    return append(data, padText...)
}

// pkcs7Unpad 去除PKCS7填充
func pkcs7Unpad(data []byte) ([]byte, error) {
    length := len(data)
    if length == 0 {
        return nil, fmt.Errorf("数据为空")
    }
    
    padding := int(data[length-1])
    if padding > length {
        return nil, fmt.Errorf("无效的填充")
    }
    
    return data[:length-padding], nil
}

// GenerateLicense 生成许可证
func GenerateLicense(hardwareHash string, expiryDateStr string, admin bool) (*License, error) {
    // 构造北京时间（东八区）23:59:59
    bjDateStr := fmt.Sprintf("%sT23:59:59+08:00", expiryDateStr)
    bjDate, err := time.Parse(time.RFC3339, bjDateStr)
    if err != nil {
        return nil, fmt.Errorf("解析过期时间失败: %v", err)
    }
    
    expiryDate := bjDate.UTC().Format(time.RFC3339)
    
    // 加密敏感数据（过期时间和管理员标识）
    encryptedExpiry, err := encryptLicenseData(expiryDate, admin, hardwareHash)
    if err != nil {
        return nil, fmt.Errorf("加密敏感数据失败: %v", err)
    }
    
    // 签名数据（包含硬件哈希和加密的敏感数据）
    licenseData := map[string]string{
        "hardwareHash":    hardwareHash,
        "encryptedExpiry": encryptedExpiry,
    }
    
    // 读取私钥
    privateKeyPath, err := getKeyFilePath("private.pem")
    if err != nil {
        return nil, fmt.Errorf("获取私钥文件路径失败: %v", err)
    }
    
    privateKeyData, err := ioutil.ReadFile(privateKeyPath)
    if err != nil {
        return nil, fmt.Errorf("读取私钥文件失败: %v", err)
    }
    
    // 解析私钥
    privateKeyBlock, _ := pem.Decode(privateKeyData)
    if privateKeyBlock == nil {
        return nil, fmt.Errorf("解析私钥失败")
    }
    
    privateKey, err := x509.ParsePKCS8PrivateKey(privateKeyBlock.Bytes)
    if err != nil {
        return nil, fmt.Errorf("解析PKCS8私钥失败: %v", err)
    }
    
    rsaPrivateKey, ok := privateKey.(*rsa.PrivateKey)
    if !ok {
        return nil, fmt.Errorf("私钥不是有效的RSA私钥")
    }
    
    // 序列化许可证数据
    licenseBytes, err := json.Marshal(licenseData)
    if err != nil {
        return nil, fmt.Errorf("序列化许可证数据失败: %v", err)
    }
    
    // 计算SHA256哈希
    hashed := sha256.Sum256(licenseBytes)
    
    // 签名
    signature, err := rsa.SignPKCS1v15(rand.Reader, rsaPrivateKey, crypto.SHA256, hashed[:])
    if err != nil {
        return nil, fmt.Errorf("签名失败: %v", err)
    }
    
    // 创建完整许可证（不包含明文敏感信息）
    license := &License{
        HardwareHash:    hardwareHash,
        EncryptedExpiry: encryptedExpiry,
        Signature:       base64.StdEncoding.EncodeToString(signature),
    }
    
    return license, nil
}

// VerifyLicense 验证许可证
func VerifyLicense() (*LicenseStatus, error) {
    // 检查license.json是否存在
    if _, err := os.Stat("license.json"); os.IsNotExist(err) {
        return &LicenseStatus{
            Valid:  false,
            Reason: "未找到license.json",
        }, nil
    }
    
    // 读取license.json
    licenseData, err := ioutil.ReadFile("license.json")
    if err != nil {
        return nil, fmt.Errorf("读取license.json失败: %v", err)
    }
    
    var license License
    if err := json.Unmarshal(licenseData, &license); err != nil {
        return nil, fmt.Errorf("解析license.json失败: %v", err)
    }
    
    // 1. 验证硬件指纹
    localHash := GetHardwareHash("")
    if license.HardwareHash != localHash {
        return &LicenseStatus{
            Valid:  false,
            Reason: "硬件指纹不匹配",
        }, nil
    }
    
    // 2. 验证签名
    publicKeyPath, err := getKeyFilePath("public.pem")
    if err != nil {
        return nil, fmt.Errorf("获取公钥文件路径失败: %v", err)
    }
    
    publicKeyData, err := ioutil.ReadFile(publicKeyPath)
    if err != nil {
        return nil, fmt.Errorf("读取公钥文件失败: %v", err)
    }
    
    // 解析公钥
    publicKeyBlock, _ := pem.Decode(publicKeyData)
    if publicKeyBlock == nil {
        return nil, fmt.Errorf("解析公钥失败")
    }
    
    publicKey, err := x509.ParsePKIXPublicKey(publicKeyBlock.Bytes)
    if err != nil {
        return nil, fmt.Errorf("解析公钥失败: %v", err)
    }
    
    rsaPublicKey, ok := publicKey.(*rsa.PublicKey)
    if !ok {
        return nil, fmt.Errorf("公钥不是有效的RSA公钥")
    }
    
    // 准备验证数据
    verifyData := map[string]string{
        "hardwareHash":    license.HardwareHash,
        "encryptedExpiry": license.EncryptedExpiry,
    }
    
    verifyBytes, err := json.Marshal(verifyData)
    if err != nil {
        return nil, fmt.Errorf("序列化验证数据失败: %v", err)
    }
    
    // 计算SHA256哈希
    hashed := sha256.Sum256(verifyBytes)
    
    // 解码签名
    signature, err := base64.StdEncoding.DecodeString(license.Signature)
    if err != nil {
        return &LicenseStatus{
            Valid:  false,
            Reason: "许可证签名格式无效",
        }, nil
    }
    
    // 验证签名
    err = rsa.VerifyPKCS1v15(rsaPublicKey, crypto.SHA256, hashed[:], signature)
    if err != nil {
        return &LicenseStatus{
            Valid:  false,
            Reason: "许可证签名无效",
        }, nil
    }
    
    // 3. 解密并验证敏感数据
    sensitiveData, err := decryptLicenseData(license.EncryptedExpiry, license.HardwareHash)
    if err != nil {
        return &LicenseStatus{
            Valid:  false,
            Reason: "许可证格式错误或已损坏",
        }, nil
    }
    
    // 如果是管理员账号，直接通过验证
    if sensitiveData.Admin {
        return &LicenseStatus{
            Valid:  true,
            Reason: "管理员账号",
        }, nil
    }
    
    // 4. 验证过期时间
    expiryDate, err := time.Parse(time.RFC3339, sensitiveData.ExpiryDate)
    if err != nil {
        return nil, fmt.Errorf("解析过期时间失败: %v", err)
    }
    
    now := time.Now().UTC()
    // 转换为北京时间进行比较（UTC+8）
    nowBJT := now.Add(8 * time.Hour)
    
    if nowBJT.After(expiryDate) {
        return &LicenseStatus{
            Valid:  false,
            Reason: "许可证已过期",
        }, nil
    }
    
    return &LicenseStatus{
        Valid:  true,
        Reason: "许可证校验通过",
    }, nil
}

func GenerateHardwareHash(CPUID string, CPUName string, MacAddress string, publicIP string) string {
    raw := fmt.Sprintf("%s-%s-%s", MacAddress, CPUID, CPUName)
    if publicIP != "" {
        raw += fmt.Sprintf("-%s", publicIP)
    }
    
    return calculateSHA256Hash(raw)
}

// GetHardwareHash 获取硬件指纹
func GetHardwareHash(publicIP string) string {
    info, err := hardware.GetHardwareInfo()
    if err != nil {
        return ""
    }
    
    raw := fmt.Sprintf("%s-%s-%s", info.MacAddress, info.CPUID, info.CPUName)
    if publicIP != "" {
        raw += fmt.Sprintf("-%s", publicIP)
    }
    
    return calculateSHA256Hash(raw)
}
